import sys
sys.path.append("/home/bachelor/zhangyi/my_code/vasp")

import os
from pymatgen.io.cif import CifParser  
from pymatgen.io.vasp import Poscar
from convert import outRoot
from utils import getAllFileAbsPathInFolder


def convert(file, savePath):
    """
    将cif转为POSCAR,用于VASP计算

    Args:
        file (str): cif文件路径
        savePath (str): 要保存的目录
    """
    cifName = file.split("/")[-1].split(".")[0]
    finalSavePath = f"{savePath}/{cifName}"
    os.makedirs(finalSavePath, exist_ok=True)
    
    # 读取CIF文件
    parser = CifParser(file)  
    structure = parser.get_structures()[0]
    
    # 转换为POSCAR格式
    poscar = Poscar(structure)
  
    # 写入POSCAR文件  
    poscar.write_file(f"{finalSavePath}/POSCAR")
    
    # 获取POTCAR文件
    os.chdir(finalSavePath)
    os.system("vaspkit -task 103")

    # 获取KPOINTS文件
    os.system("vaspkit -task 102 2 0.03")

    # 复制INCAR
    os.system("cp /home/bachelor/zhangyi/my_code/vasp/input/INCAR ./ ")


def batchProcessCif(path):
    """
    批处理CIF文件

    Args:
        path (str): CIF根路径,改目录下包含一系列.cif
    """
    cifType = path.split("/")[-1]
    cofList = getAllFileAbsPathInFolder(path)
    os.makedirs(f"{outRoot}/{cifType}", exist_ok=True)
    for cof in cofList:
        if cof.endswith(".cif"):
            convert(cof, f"{outRoot}/{cifType}")


if __name__ == '__main__':
    # 执行core-cof转换
    batchProcessCif("/home/bachelor/zhangyi/my_code/vasp/input/CORECOF-1242")